import requests
import json
import time
import traceback
from typing import Dict, List, Union, Optional

# 兼容两种 logger 导入方式
try:
    from bisect_log_config import logger
except ImportError:
    try:
        from log_config import logger
    except ImportError:
        import logging
        logger = logging.getLogger(__name__)

class ManticoreClient:
    def __init__(self, host="localhost", port=9308):
        self.base_url = f"http://{host}:{port}"

    def insert(self, index: str, id: int, document: dict) -> bool:
        """Insert document"""
        # 防御性检查：确保不包含 id 字段
        if "id" in document:
            logger.warning("Document contains redundant id field, removing")
            document = {k: v for k, v in document.items() if k != "id"}
            
        result = self._request("insert", index, id, document)
        
        return result
    
    def replace(self, index: str, id: int, document: dict) -> bool:
        """替换文档""" 
        return self._request("replace", index, id, document)

    def update(self, index: str, id: int, document: dict) -> bool:
        """部分更新文档（仅修改指定字段）"""
        return self._request("update", index, id, document)

    def search(self,
               index: str,
               query: dict,
               limit: int = 100,
               options: Optional[dict] = None,
               sort: Optional[list] = None) -> Optional[List[Dict]]:
        """
        Manticore 标准 SQL 风格搜索方法
        """
        try:
            # 参数验证
            if not index or not query:
                logger.error("无效的搜索参数")
                return None
                
            # 防御性类型转换
            if not isinstance(limit, int) or limit <= 0:
                logger.warning(f"无效的limit值: {limit}, 使用默认值100")
                limit = 100
                

            request_body = {
                "index": index,
                "limit": limit,
                "query": query
            }
            if options:
                request_body["options"] = options
            if sort:
                request_body["sort"] = sort

            resp = requests.post(
                f"{self.base_url}/search",
                json=request_body,
                timeout=30
            )

            if resp.status_code != 200:
                logger.error(f"Search request failed with status code {resp.status_code}: {resp.text}")
                logger.error(f"Failed request was: {json.dumps(request_body, indent=2, ensure_ascii=False)}")
                return None

            result = resp.json()
            
            
            # 处理两种可能的响应格式
            if 'hits' in result and 'hits' in result['hits']:
                return [
                    {**hit.get('_source', {}), 'id': hit.get('_id')}
                    for hit in result['hits']['hits']
                ]
            elif 'data' in result:
                return result['data']
            else:
                # 直接数组格式或其他格式
                logger.debug(f"Unknown response format: {list(result.keys())}")
                logger.debug(f"DEBUG - 完整未知响应示例: {json.dumps(result, indent=2, ensure_ascii=False)[:500]}")
                return result if isinstance(result, list) else []

        except requests.exceptions.RequestException:
            return None

    def replace_with_retry(self, index: str, id: int, document: dict, retries: int = 3) -> bool:
        """带重试机制的替换操作"""
        for i in range(retries):
            if self.replace(index, id, document):
                return True
            time.sleep(2 ** i)
        return False

    def batch_insert(self, index: str, documents: Union[List[dict], Dict[int, dict]]) -> bool:
        """
        批量插入文档

        Args:
            index: 索引名称
            documents: 文档列表或字典
                      - List[dict]: 自动生成ID的文档列表
                      - Dict[int, dict]: 指定ID的文档字典

        Returns:
            bool: 是否成功

        Examples:
            # 自动生成ID
            client.batch_insert("bisect", [
                {"field1": "value1", "field2": "value2"},
                {"field1": "value3", "field2": "value4"}
            ])

            # 指定ID
            client.batch_insert("bisect", {
                1001: {"field1": "value1", "field2": "value2"},
                1002: {"field1": "value3", "field2": "value4"}
            })
        """
        try:
            bulk_body = []

            if isinstance(documents, list):
                # 列表格式：自动生成ID
                for doc in documents:
                    # 确保不包含 id 字段
                    if "id" in doc:
                        doc = {k: v for k, v in doc.items() if k != "id"}
                    # 使用 insert 操作，让 Manticore 自动生成 ID
                    bulk_body.append(json.dumps({"insert": {"_index": index}}))
                    bulk_body.append(json.dumps(doc))
            elif isinstance(documents, dict):
                # 字典格式：使用指定的ID
                for doc_id, doc in documents.items():
                    # 确保不包含 id 字段
                    if "id" in doc:
                        doc = {k: v for k, v in doc.items() if k != "id"}
                    bulk_body.append(json.dumps({"insert": {"_index": index, "_id": doc_id}}))
                    bulk_body.append(json.dumps(doc))
            else:
                logger.error(f"不支持的文档格式: {type(documents)}")
                return False

            if not bulk_body:
                logger.warning("没有要插入的文档")
                return True

            # 添加结尾换行符（NDJSON格式要求）
            bulk_data = "\n".join(bulk_body) + "\n"

            logger.debug(f"批量插入 | 索引={index}, 文档数={len(bulk_body)//2}")
            logger.debug(f"请求体前500字符: {bulk_data[:500]}")

            resp = requests.post(
                f"{self.base_url}/bulk",
                data=bulk_data,
                headers={"Content-Type": "application/x-ndjson"},
                timeout=30
            )

            logger.debug(f"批量插入响应 | 状态码={resp.status_code}")

            if resp.status_code != 200:
                logger.error(f"批量插入失败 | 状态码={resp.status_code}, 响应={resp.text[:500]}")
                return False

            # 解析响应检查是否有错误
            try:
                result = resp.json()
                if "errors" in result and result["errors"]:
                    logger.error(f"批量插入有错误: {result}")
                    return False
                logger.info(f"成功批量插入 {len(bulk_body)//2} 个文档到索引 {index}")
            except json.JSONDecodeError:
                # 有些版本的 Manticore 可能不返回 JSON
                logger.debug("批量插入完成（无法解析响应JSON）")

            return True

        except requests.exceptions.RequestException as e:
            logger.error(f"批量插入异常 | 类型={type(e).__name__}, 错误={str(e)}")
            logger.debug(f"异常堆栈:\n{traceback.format_exc()}")
            return False
        except Exception as e:
            logger.error(f"批量插入处理异常: {str(e)}")
            logger.debug(f"异常堆栈:\n{traceback.format_exc()}")
            return False

    def bulk_replace(self, index: str, documents: Dict[int, dict]) -> bool:
        """批量替换文档"""
        try:
            bulk_body = []
            for doc_id, doc in documents.items():
                bulk_body.append(json.dumps({"replace": {"_index": index, "_id": doc_id}}))
                bulk_body.append(json.dumps(doc))

            resp = requests.post(
                f"{self.base_url}/bulk",
                data="\n".join(bulk_body),
                headers={"Content-Type": "application/x-ndjson"},
                timeout=10
            )
            return resp.status_code == 200
        except requests.exceptions.RequestException as e:
            # DEBUG: 捕获并打印完整异常
            logger.debug(f"DEBUG - 搜索异常 | 类型={type(e).__name__}, 错误={str(e)}")
            logger.debug(f"DEBUG - 异常堆栈:\n{traceback.format_exc()}")
            return False

    def _request(self, endpoint: str, index: str, id: int, doc: dict) -> bool:
        try:
            # 防御性检查：确保不包含 id 字段
            if "id" in doc:
                logger.error("⚠️ 非法操作：文档包含 id 字段（主键不可更新）")
                logger.error(f"文档内容: {json.dumps(doc, indent=2)}")
                doc = {k: v for k, v in doc.items() if k != "id"}
                
            # 收集调用栈信息
            from inspect import currentframe, getouterframes
            caller_info = []
            frames = getouterframes(currentframe())
            for frame in frames[1:6]:  # 获取最近的5层调用栈
                caller_info.append(f"{frame.filename}:{frame.lineno} ({frame.function})")
            
            # 添加防御性检查
            if 'j' in doc and doc['j'] is None:
                logger.error("检测到无效的 null j 字段，已自动清理")
                doc = {k: v for k, v in doc.items() if k != 'j'}

            url = f"{self.base_url}/{endpoint}"
            payload = {
                "index": index,
                "id": id,
                "doc": doc
            }
            
            resp = requests.post(
                url,
                json=payload,
                timeout=10
            )
            
            
            if resp.status_code != 200:
                # 详细错误日志
                logger.error(f"请求失败 | 状态码: {resp.status_code}")
                logger.error(f"请求端点: {endpoint}")
                logger.error(f"索引: {index}")
                logger.error(f"文档ID: {id}")
                # 检查问题字段
                if 'j' in doc:
                    logger.error(f"字段 'j' 的值类型: {type(doc['j']).__name__}")
                    logger.error(f"字段 'j' 的内容片段: {str(doc['j'])[:100]}")
                else:
                    logger.error("文档中不存在 'j' 字段")
                
                logger.error(f"完整响应: {resp.text}")
            
            return resp.status_code == 200
        except requests.exceptions.RequestException as e:
            # 异常时的调用栈信息
            logger.error(f"请求异常 | 类型: {type(e).__name__}, 错误: {str(e)}")
            logger.error(f"异常堆栈:\n{traceback.format_exc()}")
            return False

    def sql_select(self, sql_query: str) -> Optional[List[Dict]]:
        """
        执行 SQL SELECT 查询

        Args:
            sql_query: SQL 查询语句

        Returns:
            查询结果列表，失败返回 None
        """
        try:
            logger.debug(f"执行 SQL 查询: {sql_query[:200]}...")

            # 使用 text/plain 格式发送 SQL 查询（更直接，兼容性更好）
            resp = requests.post(
                f"{self.base_url}/sql",
                data=sql_query,
                headers={"Content-Type": "text/plain"},
                timeout=30
            )

            if resp.status_code != 200:
                logger.error(f"SQL查询失败 | 状态码: {resp.status_code} | 响应: {resp.text}")
                return None

            result = resp.json()
            logger.debug(f"SQL 响应格式: {type(result)} | 顶层字段: {list(result.keys()) if isinstance(result, dict) else 'list'}")

            # ManticoreSearch SQL返回的可能是搜索格式，需要解析hits结构
            if 'hits' in result and 'hits' in result['hits']:
                # 搜索格式：提取_source字段，并将_id映射为id
                data = []
                for hit in result['hits']['hits']:
                    # 对于聚合查询（COUNT, SUM等），结果在hit对象顶层，不在_source中
                    row = {}

                    # 首先提取 _source 中的字段
                    if '_source' in hit:
                        row = hit['_source'].copy()

                    # 然后添加聚合字段（count(*), sum(*), avg(*) 等）
                    # 这些字段直接在 hit 对象中，不在 _source 里
                    for key, value in hit.items():
                        if key not in ['_id', '_score', '_source', '_index']:
                            row[key] = value

                    # 将ManticoreSearch的_id映射为标准的id字段
                    if '_id' in hit:
                        row['id'] = hit['_id']

                    data.append(row)
                logger.debug(f"从 hits 格式提取了 {len(data)} 行数据")
                if data and len(data) > 0:
                    logger.debug(f"第一行字段: {list(data[0].keys())}")
                return data
            elif 'data' in result:
                # SQL格式：直接返回data
                logger.debug(f"从 data 字段提取了 {len(result['data'])} 行数据")
                return result['data']
            elif isinstance(result, list) and len(result) > 0 and 'data' in result[0]:
                # Manticore SQL 响应格式: [{"data": [...], "columns": [...], "total": N, "error": ""}]
                if 'error' in result[0] and result[0]['error']:
                    logger.error(f"SQL 查询错误: {result[0]['error']}")
                    return None
                return result[0]['data']
            else:
                logger.warning(f"未预期的响应格式: {list(result.keys()) if isinstance(result, dict) else 'unknown'}")
                return result if isinstance(result, list) else []

        except requests.exceptions.RequestException as e:
            logger.error(f"SQL 查询异常: {str(e)}")
            return None
        except Exception as e:
            logger.error(f"SQL 查询处理异常: {str(e)}")
            logger.error(traceback.format_exc())
            return None

    def sql_raw(self, sql_query: str) -> Optional[List[Dict]]:
        """
        执行原始 SQL 查询（UPDATE, DELETE 等）

        使用 /sql?mode=raw 端点，支持所有类型的 SQL 语句

        Args:
            sql_query: SQL 语句

        Returns:
            查询结果，格式为 [{"total": N, "error": "", "warning": ""}]
        """
        try:
            logger.debug(f"执行原始 SQL: {sql_query[:200]}...")

            # 使用 /sql?mode=raw 端点支持 UPDATE/DELETE 等非 SELECT 查询
            resp = requests.post(
                f"{self.base_url}/sql?mode=raw",
                data=sql_query,
                headers={"Content-Type": "text/plain"},
                timeout=30
            )

            if resp.status_code != 200:
                logger.error(f"原始SQL执行失败 | 状态码: {resp.status_code} | 响应: {resp.text}")
                return [{"error": f"HTTP {resp.status_code}: {resp.text}"}]

            result = resp.json()

            # 返回原始结果（确保返回 List[Dict] 格式）
            return result if isinstance(result, list) else [result]

        except requests.exceptions.RequestException as e:
            logger.error(f"原始SQL异常: {str(e)}")
            return [{"error": str(e)}]
        except Exception as e:
            logger.error(f"原始SQL处理异常: {str(e)}")
            logger.error(traceback.format_exc())
            return [{"error": str(e)}]
